import * as tf from "@tensorflow/tfjs";
import {
  CLASSES,
  NUM_CHANNELS,
  IMAGE_WIDTH,
  IMAGE_HEIGHT,
  IMAGE_SIZE,
  NUM_CLASSES,
  NUM_DATASET_ELEMENTS,
  NUM_TRAIN_ELEMENTS,
  NUM_TEST_ELEMENTS,
  saveDir,
  imagesFileName,
  labelsFileName,
} from "../datasets/OCRDatasetConstants";
import { getFile } from "@/renderer/api/messageAPI";

export default class OCRData {
  classes = CLASSES
  imageWidth = IMAGE_WIDTH;
  imageHeight = IMAGE_HEIGHT;
  imageSize = IMAGE_SIZE;
  numClasses = NUM_CLASSES;
  numDatasetElements = NUM_DATASET_ELEMENTS;
  numTrainElements = NUM_TRAIN_ELEMENTS;
  numTestElements = NUM_TEST_ELEMENTS;
  datasetImages;
  datasetLabels;
  trainImages;
  testImages;
  trainLabels;
  testLabels;
  trainIndex = 0;
  testIndex = 0;

  constructor() {}

  async load() {
    const imgRequest = getFile(saveDir + imagesFileName);
    const labelsRequest = getFile(saveDir + labelsFileName);

    const [datasetImages, datasetLabels] = await Promise.all([
      imgRequest,
      labelsRequest,
    ]);

    this.datasetImages = datasetImages
    this.datasetLabels = new Uint32Array(datasetLabels.buffer)

    this.trainImages = this.datasetImages.slice(0, IMAGE_SIZE * NUM_TRAIN_ELEMENTS * NUM_CHANNELS)
    this.testImages = this.datasetImages.slice(IMAGE_SIZE * NUM_TRAIN_ELEMENTS * NUM_CHANNELS)

    this.trainLabels = this.datasetLabels.slice(0, NUM_TRAIN_ELEMENTS);
    this.testLabels = this.datasetLabels.slice(NUM_TRAIN_ELEMENTS);
  }

  static imagesDataHandle(imagesData: Uint8Array, numElements) {
    // 单通道4字节buffer
    const datasetBytesBuffer = new ArrayBuffer(
      numElements * IMAGE_SIZE * 4
    );
    const datasetImages = new Float32Array(datasetBytesBuffer);
    // 归一化
    for (let i = 0; i < imagesData.length / 4; i++) {
      datasetImages[i] = imagesData[i * 4] / 255;
    }

    const xs = tf.tensor4d(datasetImages, [
      numElements,
      IMAGE_HEIGHT,
      IMAGE_WIDTH,
      1,
    ]);

    return xs
  }
  
  static labelsDataHandle(labelsData: Uint32Array, numElements) {
    // oneHot有尺寸限制
    // 独热表示法
    // const datasetLabels = tf
    // .tidy(() => {
    //   return tf.oneHot(tf.tensor1d(Int32Array.from(labelsData), "int32"), NUM_CLASSES);
    // })
    // .dataSync();

    // const labels = tf.tensor2d(datasetLabels, [
    //   numElements,
    //   NUM_CLASSES,
    // ]);

    // 稀疏表示法（索引表示法）
    const labels = tf.tensor1d(Float32Array.from(labelsData))

    return labels
  }

  getTrainData(numElements:number = NUM_TRAIN_ELEMENTS) {
    const index = this.trainIndex

    const trainImages = this.trainImages.slice(index, index + numElements * IMAGE_SIZE * NUM_CHANNELS)
    const trainLabels = this.trainLabels.slice(index, index + numElements)

    const realNumElements = trainLabels.length // slice去掉超出数量

    const xs = OCRData.imagesDataHandle(trainImages, realNumElements)
    const labels = OCRData.labelsDataHandle(trainLabels, realNumElements)

    this.trainIndex += realNumElements

    return { xs, labels };
  }

  getTestData(numElements: number = NUM_TEST_ELEMENTS) {
    const index = this.testIndex

    const testImages = this.testImages.slice(index, index + numElements * IMAGE_SIZE * NUM_CHANNELS)
    const testLabels = this.testLabels.slice(index, index + numElements)

    const realNumElements = testLabels.length

    const xs = OCRData.imagesDataHandle(testImages, realNumElements)
    const labels = OCRData.labelsDataHandle(testLabels, realNumElements)

    this.testIndex += realNumElements
    
    return { xs, labels };
  }

  getTrainDataset(batchSize = 32) {
    const imagesData = this.trainImages
    const datasetBytesBuffer = new ArrayBuffer(
      NUM_TRAIN_ELEMENTS * IMAGE_SIZE * 4
    );
    const trainImages: any = new Float32Array(datasetBytesBuffer);

    for (let i = 0; i < imagesData.length / 4; i++) {
      trainImages[i] = imagesData[i * 4] / 255;
    }

    const xs: any = tf.tensor4d(trainImages, [
      NUM_TRAIN_ELEMENTS,
      IMAGE_HEIGHT,
      IMAGE_WIDTH,
      1,
    ]);

    const trainImagesDataset = tf.data.array(xs.arraySync())
    const trainLabelsDataset = tf.data.array(this.trainLabels) // 索引表示法
    const trainDataset = tf.data.zip({ xs: trainImagesDataset, ys: trainLabelsDataset }).batch(batchSize)

    return trainDataset
  }

  getTestDataset(batchSize = 32) {
    const imagesData = this.testImages
    const datasetBytesBuffer = new ArrayBuffer(
      NUM_TEST_ELEMENTS * IMAGE_SIZE * 4
    );
    const testImages: any = new Float32Array(datasetBytesBuffer);

    for (let i = 0; i < imagesData.length / 4; i++) {
      testImages[i] = imagesData[i * 4] / 255;
    }

    const xs: any = tf.tensor4d(testImages, [
      NUM_TEST_ELEMENTS,
      IMAGE_HEIGHT,
      IMAGE_WIDTH,
      1,
    ]);

    const testImagesDataset = tf.data.array(xs.arraySync())
    const testLabelsDataset = tf.data.array(this.testLabels)
    const testDataset = tf.data.zip({ xs: testImagesDataset, ys: testLabelsDataset }).batch(batchSize)

    return testDataset
  }
}
